import os
import shutil
from sklearn.model_selection import train_test_split

# 数据集根目录
dataset_dir = "data"
output_dir = "animals10"

# 创建输出目录
os.makedirs(os.path.join(output_dir, "train"), exist_ok=True)
os.makedirs(os.path.join(output_dir, "val"), exist_ok=True)
os.makedirs(os.path.join(output_dir, "test"), exist_ok=True)

# 遍历每个类别文件夹
for label in os.listdir(dataset_dir):
    label_dir = os.path.join(dataset_dir, label)
    if not os.path.isdir(label_dir):
        print(f"Skipping non-directory: {label_dir}")
        continue

    # 获取文件列表
    files = [f for f in os.listdir(label_dir) if f.endswith((".jpg", ".png","jpeg")) and not f.startswith('.')]
    if not files:
        print(f"No image files found in folder: {label_dir}. Skipping...")
        continue

    # 划分数据集
    train_files, temp_files = train_test_split(files, test_size=0.3, random_state=42)
    val_files, test_files = train_test_split(temp_files, test_size=0.5, random_state=42)

    # 创建目标文件夹
    for split in ["train", "val", "test"]:
        os.makedirs(os.path.join(output_dir, split, label), exist_ok=True)

    # 复制文件到目标文件夹
    for file in train_files:
        shutil.copy(os.path.join(label_dir, file), os.path.join(output_dir, "train", label, file))
    for file in val_files:
        shutil.copy(os.path.join(label_dir, file), os.path.join(output_dir, "val", label, file))
    for file in test_files:
        shutil.copy(os.path.join(label_dir, file), os.path.join(output_dir, "test", label, file))